import math

import numpy as np
import load_data
import math

MILLISECONDS_IN_SECOND = 1000.0
# B_IN_MB = 1000000.0
BITS_IN_BYTE = 8.0
RANDOM_SEED = 30
VIDEO_CHUNCK_LEN = 1000.0  # millisec, every time add this amount to buffer
BITRATE_LEVELS = 4
TOTAL_VIDEO_CHUNCK = 200
BUFFER_THRESH = 10.0 * MILLISECONDS_IN_SECOND  # millisec, max buffer limit
DRAIN_BUFFER_SLEEP_TIME = 500.0  # millisec
PACKET_PAYLOAD_PORTION = 0.95
LINK_RTT = 80  # millisec
PACKET_SIZE = 1500  # bytes
NOISE_LOW = 0.9
NOISE_HIGH = 1.1


class Environment:
    def __init__(self, all_cooked_time, all_cooked_bw, all_FoV_time, all_FoV_pos, random_seed=RANDOM_SEED):
        assert len(all_cooked_time) == len(all_cooked_bw)
        assert len(all_FoV_time) == len(all_FoV_pos)

        np.random.seed(random_seed)
        self.agent_id = random_seed
        self.all_cooked_time = all_cooked_time
        self.all_cooked_bw = all_cooked_bw
        self.all_FoV_time = all_FoV_time
        self.all_FoV_pos = all_FoV_pos

        self.video_chunk_counter = 0
        self.all_video_chunk_counter = [0, 0, 0, 0, 0, 0,
                                        0, 0, 0, 0, 0, 0,
                                        0, 0, 0, 0, 0, 0,
                                        0, 0, 0, 0, 0, 0]
        self.buffer_size = 0
        self.all_buffer_size = [0, 0, 0, 0, 0, 0,
                                0, 0, 0, 0, 0, 0,
                                0, 0, 0, 0, 0, 0,
                                0, 0, 0, 0, 0, 0]

        # pick a random trace file
        self.trace_idx = np.random.randint(len(self.all_cooked_time))
        self.cooked_time = self.all_cooked_time[self.trace_idx]
        self.cooked_bw = self.all_cooked_bw[self.trace_idx]
        self.cooked_fov = self.all_FoV_pos[self.trace_idx]

        # randomize the start point of the trace
        # note: trace file starts with time 0
        self.mahimahi_ptr = np.random.randint(1, len(self.cooked_bw))
        self.all_mahimahi_ptr = np.repeat(np.expand_dims(self.mahimahi_ptr, axis=0), 24, axis=0)

        self.last_mahimahi_time = self.cooked_time[self.mahimahi_ptr - 1]
        self.all_last_mahimahi_time = np.repeat(np.expand_dims(self.last_mahimahi_time, axis=0), 24, axis=0)

        self.fov_time = self.mahimahi_ptr
        self.all_fov_time = np.repeat(np.expand_dims(self.fov_time, axis=0), 24, axis=0)

        # self.video_size = load_data.load_segment_size('tiledSkiing')
        self.video_size = load_data.load_multi_segment_size('tiledSkiing')

    def get_video_chunk(self, tile):
        """
        模拟每个tile的下载
        :param tile:瓦片对象
        :return:
        """
        quality = tile.quality
        face_coord = tile.face_coord
        tile_id = tile.tile_id
        assert quality >= 0
        assert quality < BITRATE_LEVELS

        face = tile_id // 4
        tile_in_face = tile_id % 4
        video_chunk_size = self.video_size[self.all_video_chunk_counter[tile_id]][face][tile_in_face][quality]

        # use the delivery opportunity in mahimahi
        delay = 0.0  # in ms
        video_chunk_counter_sent = 0  # in bytes

        # tile偏离角度
        fov_tmp = self.cooked_fov[self.all_fov_time[tile_id]]
        deviation = math.acos((fov_tmp[0] * face_coord[0] + fov_tmp[1] * face_coord[1] + fov_tmp[2] * face_coord[2]) /
                              math.sqrt(fov_tmp[0] * fov_tmp[0] + fov_tmp[1] * fov_tmp[1] + fov_tmp[2] * fov_tmp[2]) /
                              math.sqrt(face_coord[0] * face_coord[0] + face_coord[1] * face_coord[1] + face_coord[2] * face_coord[2]))
        deviation /= math.pi

        while True:  # download video chunk over mahimahi
            # print(str(self.agent_id) + '\t' + str(tile_id) + '\t' + str(len(self.cooked_bw)) + '\t'
            #       + str(self.all_mahimahi_ptr[tile_id]))
            throughput = self.cooked_bw[self.all_mahimahi_ptr[tile_id]][tile_id] / BITS_IN_BYTE  # Bps
            duration = self.cooked_time[self.all_mahimahi_ptr[tile_id]] - self.all_last_mahimahi_time[tile_id]    # s

            packet_payload = throughput * duration * PACKET_PAYLOAD_PORTION

            if video_chunk_counter_sent + packet_payload > video_chunk_size:
                fractional_time = (video_chunk_size - video_chunk_counter_sent) / throughput / PACKET_PAYLOAD_PORTION
                delay += fractional_time
                self.all_last_mahimahi_time[tile_id] += fractional_time
                assert (self.all_last_mahimahi_time[tile_id] <= self.cooked_time[self.all_mahimahi_ptr[tile_id]])
                break

            video_chunk_counter_sent += packet_payload
            delay += duration
            self.all_last_mahimahi_time[tile_id] = self.cooked_time[self.all_mahimahi_ptr[tile_id]]
            self.all_mahimahi_ptr[tile_id] += 1

            if self.all_mahimahi_ptr[tile_id] >= len(self.cooked_bw):
                # loop back in the beginning
                # note: trace file starts with time 0
                self.all_mahimahi_ptr[tile_id] = 1
                self.all_last_mahimahi_time[tile_id] = 0
            self.all_fov_time[tile_id] = self.all_mahimahi_ptr[tile_id]

        delay *= MILLISECONDS_IN_SECOND     # ms
        delay += LINK_RTT
        delay *= np.random.uniform(NOISE_LOW, NOISE_HIGH)

        # rebuffer time
        rebuf = np.maximum(delay - self.all_buffer_size[tile_id], 0.0)  # ms

        # update the buffer
        self.all_buffer_size[tile_id] = np.maximum(self.all_buffer_size[tile_id] - delay, 0.0)  # ms

        # add in the new chunk
        self.all_buffer_size[tile_id] += VIDEO_CHUNCK_LEN

        # sleep if buffer gets too large
        sleep_time = 0
        if self.all_buffer_size[tile_id] > BUFFER_THRESH:
            # exceed the buffer limit
            # we need to skip some network bandwidth here
            # but do not add up the delay
            drain_buffer_time = self.all_buffer_size[tile_id] - BUFFER_THRESH
            sleep_time = np.ceil(drain_buffer_time / DRAIN_BUFFER_SLEEP_TIME) * \
                         DRAIN_BUFFER_SLEEP_TIME
            self.all_buffer_size[tile_id] -= sleep_time

            while True:
                duration = self.cooked_time[self.all_mahimahi_ptr[tile_id]] \
                           - self.all_last_mahimahi_time[tile_id]
                if duration > sleep_time / MILLISECONDS_IN_SECOND:
                    self.all_last_mahimahi_time[tile_id] += sleep_time / MILLISECONDS_IN_SECOND
                    break
                sleep_time -= duration * MILLISECONDS_IN_SECOND
                self.all_last_mahimahi_time[tile_id] = self.cooked_time[self.all_mahimahi_ptr[tile_id]]
                self.all_mahimahi_ptr[tile_id] += 1

                if self.all_mahimahi_ptr[tile_id] >= len(self.cooked_bw):
                    # loop back in the beginning
                    # note: trace file starts with time 0
                    self.all_mahimahi_ptr[tile_id] = 1
                    self.all_last_mahimahi_time[tile_id] = 0
                self.all_fov_time[tile_id] = self.all_mahimahi_ptr[tile_id]

        # the "last buffer size" return to the controller
        # Note: in old version of dash the lowest buffer is 0.
        # In the new version the buffer always have at least
        # one chunk of video
        return_buffer_size = self.all_buffer_size[tile_id]

        self.all_video_chunk_counter[tile_id] += 1
        video_chunk_remain = TOTAL_VIDEO_CHUNCK - self.all_video_chunk_counter[tile_id]

        end_of_video = False
        if self.all_video_chunk_counter[tile_id] >= TOTAL_VIDEO_CHUNCK:
            end_of_video = True
            self.all_buffer_size[tile_id] = 0
            self.all_video_chunk_counter[tile_id] = 0

            # pick a random trace file
            self.trace_idx = np.random.randint(len(self.all_cooked_time))
            self.cooked_time = self.all_cooked_time[self.trace_idx]
            self.cooked_bw = self.all_cooked_bw[self.trace_idx]
            self.cooked_fov = self.all_FoV_pos[self.trace_idx]

            # randomize the start point of the video
            # note: trace file starts with time 0
            self.mahimahi_ptr = np.random.randint(1, len(self.cooked_bw))
            self.all_mahimahi_ptr = np.repeat(np.expand_dims(self.mahimahi_ptr, axis=0), 24, axis=0)

            self.last_mahimahi_time = self.cooked_time[self.mahimahi_ptr - 1]
            self.all_last_mahimahi_time = np.repeat(np.expand_dims(self.last_mahimahi_time, axis=0), 24, axis=0)

            self.fov_time = self.mahimahi_ptr
            self.all_fov_time = np.repeat(np.expand_dims(self.fov_time, axis=0), 24, axis=0)

        next_video_chunk_sizes = []
        for i in range(BITRATE_LEVELS):
            next_video_chunk_sizes.append(round(self.video_size[self.all_video_chunk_counter[tile_id]][face][tile_in_face][i] / 1000, 3))

        # delay (ms)
        # deviation
        # sleep_time (ms)
        # return_buffer_size / MILLISECONDS_IN_SECOND (s)
        # rebuf / MILLISECONDS_IN_SECOND (s)
        # video_chunk_size (kByte)
        # next_video_chunk_sizes (kByte)
        # end_of_video (Boolean)
        # video_chunk_remain
        return delay, \
               deviation, \
               sleep_time, \
               return_buffer_size / MILLISECONDS_IN_SECOND, \
               rebuf / MILLISECONDS_IN_SECOND, \
               round(video_chunk_size / 1000, 3), \
               next_video_chunk_sizes, \
               end_of_video, \
               video_chunk_remain, \
               self.trace_idx
